package com.jboltai.scheduler;

import cn.hutool.core.util.StrUtil;
import com.jboltai.capability.model.AIEventProcessor;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.resource.ai.AIModel;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.resource.ai.AIResourceCenter;
import com.jboltai.resource.ai.AIResourceGroup;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Iterator;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * AI事件调度器
 */
class AIEventScheduler extends Scheduler<AIEvent>{

    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);


    /**
     * 每个大模型的事件队列
     */
    private  final LinkedBlockingQueue<AIEvent> eventQueue = new LinkedBlockingQueue(JBoltAIConfig.getMaxAIEventWaitCount());

    public AIEventScheduler() {
        super("ai");
    }
    /**
     * 发布一个事件
     * @param event
     */
    public void publish(AIEvent event) {
        event.setState(EventState.WAITING);
        if (event.getModel() != null) {
            boolean hasResource = AIResourceCenter.checkHasResource(event.getModel(), event.getAIResource());
            if (!hasResource) {
                //没有可用的资源
                if (event.getAIResource() == null) {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到可用的" + event.getModel() + "资源");
                    LOGGER.error("未找到可用的{}资源", event.getModel());
                } else {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getAIResource()+"的" + event.getModel() + "资源");
                    LOGGER.error("未找到id为{}的{}资源", event.getAIResource(),event.getModel());
                }
                return;
            }
        } else if (StrUtil.isNotBlank(event.getAIResource())) {
            if (AIResourceCenter.getResourceById(event.getAIResource()) ==  null) {
                event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getAIResource()+"的资源");
                LOGGER.error("未找到id为{}的资源", event.getAIResource());
                return ;
            }

        }


        boolean insert = eventQueue.offer(event);

        if (!insert) {
            //队列满了
            event.fail(EventErrorType.QUEUE_FULL);
        }
    }

    @Override
    void process() {


        try {
            //检测该模型的事件队列 是不是空的
//            if (eventQueue.isEmpty()) {
//                sleep(20);
//                continue;
//            }

            Iterator<AIEvent> iterator = eventQueue.iterator();
            while (iterator.hasNext()) {
                AIEvent event = iterator.next();

                if (event.getState() == EventState.CANCEL) {
                    //事件被取消了
                    iterator.remove();
                    continue;
                }

                dispatch(event);

                if (event.getState() != EventState.WAITING) {
                    //不是等待状态，就不用再处理了
                    iterator.remove();
                }
            }
//            sleep(10);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }



    }

    @Override
    void dispatch(AIEvent event) {
        AIResourceGroup resources = null;
        AIResource resource = null;


        if (StrUtil.isNotBlank(event.getAIResource())) {
            //事件指定了某一个资源
            resource = AIResourceCenter.getResourceById(event.getAIResource());
            if (resource ==  null) {
                event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getAIResource()+"的资源");
                LOGGER.error("未找到id为{}的资源", event.getAIResource());
                return ;
            }

        }
        if (event.getModel() != null && resource == null) {
            //事件指定了 大模型, 根据大模型找资源
            resources = AIResourceCenter.getModelResources(event.getModel());
            if (resources == null) {
                //没有可用的资源
                event.fail(EventErrorType.NO_RESOURCE, "未找到可用的" + event.getModel() + "资源");
                LOGGER.error("未找到可用的{}资源", event.getModel());
                return;
            }
            if (event.getModel().getIsLocal() && event.getModelName() != null) {
                //如果是本地部署的大模型，优先需要根据类型和默认模型名找资源，否则可能引发bug：两个ollama跑的不同模型，根据modeltype+modelName找资源a时，会匹配到另一个资源b
                resource = resources.getResourceByModelName(event.getModelName());
            }

            if (resource == null){
                resource = resources.getResource(event.getAIResource());
            }




        } else if (resource == null) {
            //事件未指定大模型，从可用资源中取一个
            resource = AIResourceCenter.getValidResource(event);
            if (resource != null) {
                resources =resource.getResourceGroup();
            }

        }

        if (event.getModel() == null && resource != null) {
            //本身没设置用哪个大模型，但是资源有了，就根据资源指定用的大模型
            event.setModel(resource.getModel());

        }
        if (event.getModelName() == null) {
            if (resource != null && resource.getDefaultModelName() != null) {
                //优先使用资源配置的大模型名称
                event.setModelName(resource.getDefaultModelName());
            } else if (event.getModel() != null){
                event.setModelName(event.getModel().getDefaultModel());
            }
        }

        int remainTimes = 0;
        if (resource == null || (remainTimes= resource.checkAvailableAndReduce()) == -1) {
            //当前没有可用的资源
            if (event.getWaitTimestamp() > 0 && System.currentTimeMillis() > event.getWaitTimestamp()) {
                //超时了还没资源可用，就算了
                event.fail(EventErrorType.QUEUE_WAIT_TIMEOUT);
                LOGGER.error("事件{} 等待超时，放弃执行", event.getId());
            }
            if (event.getState() != EventState.WAITING) {
                //不是等待状态，说明已经处理过一次了，肯定不在等待队列中了
                if (event.getAIResource() == null) {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到可用的" + event.getModel() + "资源");
                    LOGGER.error("当前没有可用的{}资源", event.getModel());
                } else {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getAIResource()+"的" + event.getModel() + "资源");
                    LOGGER.error("当前没有可用的id为{}的{}资源", event.getAIResource(),event.getModel());
                }
                return;
            }
            return;
        }
        if (remainTimes == 0) {
            //当前资源已经用完了，需要重置
            resources.update(resource);
        }
        if (event.setState(EventState.PROCESSING) == false) {
            //状态更新失败，可能已经取消了
            return;
        }
        switch (event.getType()) {
            case AI_CHAT:{
                AIEventProcessor eventProcessor = event.getModel().getAIEventProcessor();
                event.setAttr(EventAttrKey.EVENT_PROCESSOR, eventProcessor);
                event.setAttr(EventAttrKey.EVENT_RESOURCE, resource);
                eventProcessor.process((AIChatEvent) event, resource);
                break;
            }

        }

    }


}
